import { WebSocketServer, WebSocket } from "ws";
import { ChatMessage, RoomMessage, ChatRoom } from "../models/index.js";
import { User } from "../models/index.js";
import jwt from "jsonwebtoken";
import { key } from "../config/config.default.js";

const clients = new Map(); // 存储连接的客户端
const rooms = new Map(); // 存储聊天室的成员

export function initializeWebSocket(server) {
  const wss = new WebSocketServer({
    server,
  });

  wss.on("connection", async (ws, req) => {
    console.log("新的WebSocket连接");
    let userId;

    ws.on("message", async (message) => {
      try {
        const messageData = JSON.parse(message);
        const { type, token } = messageData;

        // 处理认证消息
        if (type === "auth") {
          try {
            const decoded = jwt.verify(token, key);
            userId = decoded.id;
            clients.set(userId, ws);
            console.log(`用户${userId}认证成功`);

            // 向客户端发送认证成功消息
            ws.send(
              JSON.stringify({
                type: "auth_success",
                data: { userId },
              })
            );
            return;
          } catch (error) {
            ws.send(
              JSON.stringify({
                type: "error",
                message: "认证失败",
              })
            );
            ws.close();
            return;
          }
        }

        // 如果未认证，拒绝其他类型的消息
        if (!userId) {
          ws.send(
            JSON.stringify({
              type: "error",
              message: "请先进行认证",
            })
          );
          return;
        }

        // 处理其他类型的消息
        const { roomId, content, recipientId } = messageData;
        switch (type) {
          // 加入群聊
          case "joinRoom":
            await handleJoinRoom(userId, roomId, ws);
            break;
          // 离开群聊
          case "leaveRoom":
            await handleLeaveRoom(userId, roomId);
            break;
          // 群聊消息
          case "roomMessage":
            await handleRoomMessage(userId, roomId, content);
            break;
          // 私聊消息
          case "privateMessage":
            await handlePrivateMessage(userId, recipientId, content);
            break;
          default:
            ws.send(
              JSON.stringify({
                type: "error",
                message: "未知的消息类型",
              })
            );
        }
      } catch (error) {
        console.error("消息处理错误:", error);
        ws.send(
          JSON.stringify({
            type: "error",
            message: "消息处理失败",
          })
        );
      }
    });

    ws.on("close", () => {
      if (userId) {
        console.log(`用户 ${userId} 断开连接`);
        clients.delete(userId);
        // 从所有房间中移除用户
        for (const [roomId, members] of rooms.entries()) {
          if (members.has(userId)) {
            members.delete(userId);
            broadcastToRoom(roomId, {
              type: "userLeft",
              data: { userId },
            });
          }
        }
      }
    });

    ws.on("error", (error) => {
      console.error("WebSocket错误:", error);
      if (userId) {
        clients.delete(userId);
      }
    });
  });
}

// 处理加入房间
async function handleJoinRoom(userId, roomId, ws) {
  try {
    const room = await ChatRoom.findById(roomId);
    if (!room) {
      ws.send(
        JSON.stringify({
          type: "error",
          message: "房间不存在",
        })
      );
      return;
    }

    // 将用户添加到房间成员列表
    if (!rooms.has(roomId) || !rooms.get(roomId).has(userId)) {
      rooms.set(roomId, new Set());
    }
    rooms.get(roomId).add(userId);
    // 数据库中添加用户到群聊
    await ChatRoom.updateOne(
      { _id: roomId },
      { $addToSet: { members: userId } }
    );

    // 发送用户加入群聊的系统消息
    await RoomMessage.create({
      roomId,
      senderId: "system",
      content: `${userId} 加入了群聊`,
      timestamp: new Date(),
      type: "system",
    });

    // 获取最近的消息历史
    const messages = await RoomMessage.find({ roomId })
      .sort({ timestamp: -1 })
      .limit(10)
      .populate("senderId", "username avatar");

    // 发送加入成功消息和历史消息
    ws.send(
      JSON.stringify({
        type: "joinedRoom",
        data: {
          roomId,
          messages: messages.reverse(),
        },
      })
    );

    // 广播新用户加入消息
    broadcastToRoom(roomId, {
      type: "userJoined",
      data: { userId, roomId },
    });
  } catch (error) {
    ws.send(
      JSON.stringify({
        type: "error",
        message: "加入房间失败",
      })
    );
  }
}

// 处理离开房间
async function handleLeaveRoom(userId, roomId) {
  if (rooms.has(roomId)) {
    rooms.get(roomId).delete(userId);
    broadcastToRoom(roomId, {
      type: "userLeft",
      data: { userId },
    });
  }
}

// 处理群聊消息
async function handleRoomMessage(userId, roomId, content) {
  try {
    const room = await ChatRoom.findById(roomId);
    console.log(room.members.includes(userId));

    // 检查用户是否在房间中
    if (!room.members.includes(userId)) {
      throw new Error("您不在该房间中");
    }

    // 保存消息
    const message = new RoomMessage({
      roomId,
      senderId: userId,
      content,
      timestamp: new Date(),
    });
    await message.save();

    // 获取发送者信息
    const sender = await User.findById(userId);

    // 广播消息给房间所有成员
    const messageData = {
      type: "roomMessage",
      data: {
        messageId: message._id,
        roomId,
        senderId: userId,
        senderName: sender.username,
        senderAvatar: sender.avatar,
        content,
        timestamp: message.timestamp,
      },
    };

    broadcastToRoom(roomId, messageData);
  } catch (error) {
    const ws = clients.get(userId);
    if (ws) {
      ws.send(
        JSON.stringify({
          type: "error",
          message: error.message,
        })
      );
    }
  }
}

// 广播消息给房间所有成员
function broadcastToRoom(roomId, message) {
  const members = rooms.get(roomId);
  if (members) {
    members.forEach((userId) => {
      const ws = clients.get(userId);
      if (ws && ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify(message));
      }
    });
  }
}

// 处理私聊消息
async function handlePrivateMessage(senderId, recipientId, content) {
  try {
    // 保存消息到数据库
    const message = new ChatMessage({
      senderId,
      recipientId,
      content,
      timestamp: new Date(),
      read: false,
    });
    await message.save();

    // 获取发送者信息
    const sender = await User.findById(senderId);

    // 构建消息数据
    const messageData = {
      type: "privateMessage",
      data: {
        messageId: message._id,
        senderId,
        senderName: sender.username,
        senderAvatar: sender.avatar,
        content,
        timestamp: message.timestamp,
        read: false,
      },
    };

    // 发送给接收者
    const recipientWs = clients.get(recipientId);
    if (recipientWs && recipientWs.readyState === WebSocket.OPEN) {
      recipientWs.send(JSON.stringify(messageData));
    }

    // 发送确认给发送者
    const senderWs = clients.get(senderId);
    if (senderWs && senderWs.readyState === WebSocket.OPEN) {
      senderWs.send(
        JSON.stringify({
          ...messageData,
          type: "messageSent",
        })
      );
    }
  } catch (error) {
    const ws = clients.get(senderId);
    if (ws) {
      ws.send(
        JSON.stringify({
          type: "error",
          message: error.message,
        })
      );
    }
  }
}
